#ifndef TINYNDARRAY_INTERNAL_REDUCTION_H
#define TINYNDARRAY_INTERNAL_REDUCTION_H
#include <functional>
#include <sstream>
#include <type_traits>
#include <vector>

#include "tinyndarray.h"

namespace tinyndarray {

enum class ReduceType { SUM, MAX, MIN, UNKNOWN };
// 带归约的并行执行
template <typename F, typename R>
float RunParallelWithReduce(int size, F op, R reduce, float init_v) {
    // 决定并行参数
    int n_workers = -1, n_batch = -1, batch_size = -1;
    GetParallelParams(size, n_workers, n_batch, batch_size);

    if (n_workers <= 1) {
        // 单线程执行
        float v = init_v;
        for (int i = 0; i < size; i++) {
            v = reduce(v, op(i));
        }
        return v;
    } else {
        // 并行执行
        std::atomic<int> next_batch(0);
        std::vector<std::thread> workers;
        std::vector<float> results(n_workers);
        workers.reserve(n_workers);

        for (int w = 0; w < n_workers; w++) {
            workers.emplace_back([=, &next_batch, &results, &op, &reduce]() {
                float v = init_v;
                while (true) {
                    int batch = next_batch++;
                    if (batch >= n_batch) break;

                    int start = batch * batch_size;
                    int end = std::min(start + batch_size, size);

                    for (int i = start; i < end; i++) {
                        v = reduce(v, op(i));
                    }
                }
                results[w] = v;
            });
        }

        for (auto& worker : workers) {
            worker.join();
        }

        // 归约结果
        float result = init_v;
        for (float v : results) {
            result = reduce(result, v);
        }
        return result;
    }
}

}

#endif